import sys
import os
sys.path.append(os.path.abspath("/home/xilinx/jupyter_notebooks/dac_sdc_2021/common"))

import argparse
import time
import numpy as np
import pynq

from IoU import *
import dac_sdc
import cBBox
import jpgLoader
import AOS_Tools

parser = argparse.ArgumentParser(description="MobileNetv2 Tunable Activation Imbalance Transfer Quantization")
parser.add_argument('--f', type=int, default=333)
parser.add_argument('--v', type=int, default=850)
args = parser.parse_args()
print(args)

os.system("sudo sh -c \"sync && echo 3 > /proc/sys/vm/drop_caches\"")

BATCH_SIZE = 25
IMAGE_ROW, IMAGE_COL = 360, 640
N_THREAD = 8

cLoader = jpgLoader.jpgLoader(N_THREAD, IMAGE_ROW * IMAGE_COL * 3, 0x2)
cboxer = cBBox.cBBox(BATCH_SIZE)

result = np.zeros(shape=(1000 // BATCH_SIZE, BATCH_SIZE, 4), dtype=np.int32)
out_cnt = -1
ping, pong = 0, 1

img = [None] * 2
img[0] = pynq.allocate(shape=(BATCH_SIZE, IMAGE_ROW * IMAGE_COL * 3), dtype=np.uint8, cacheable=1)
img[1] = pynq.allocate(shape=(BATCH_SIZE, IMAGE_ROW * IMAGE_COL * 3), dtype=np.uint8, cacheable=1)
out = [None] * 2
out[0] = pynq.allocate(shape=(BATCH_SIZE, 2, 7), dtype=np.int16, cacheable=1)
out[1] = pynq.allocate(shape=(BATCH_SIZE, 2, 7), dtype=np.int16, cacheable=1)
print("Allocating memory done")

bitfile = "AOS-SkyNet.bit"
overlay = pynq.Overlay(bitfile)
PulseGen = overlay.PulseGen
OSM = overlay.OSM_Array
AOS_Tools.set_frequency(0, args.f)

dma = overlay.axi_dma
print("Bitstream loaded")

total_time = 0
print("******Start******")
power_list = []
osm_list = []
for i in range(1):
	team_name = "SkrSkr"
	team = dac_sdc.Team(team_name, batch_size = BATCH_SIZE)
	img_paths = team.get_next_batch()
	path = [str(img_path) for img_path in img_paths]
	cLoader.start(BATCH_SIZE)
	cLoader.load(path, img[ping])
	cLoader.wait()
	out_cnt = -1
	ping, pong = 0, 1
	while not img_paths is None:
		img_paths = team.get_next_batch()
		if not img_paths is None:
			path = [str(img_path) for img_path in img_paths]
			cLoader.start(BATCH_SIZE)
			cLoader.load(path, img[pong])
		start = time.time()
		dma.sendchannel.transfer(img[ping])
		dma.recvchannel.transfer(out[ping])
		power_list.append(AOS_Tools.ultra96v2_read_power(3))
		dma.recvchannel.wait()
		dma.sendchannel.wait()
		end = time.time()
		total_time += end - start

		cLoader.wait()
		cboxer.compute(out[pong], result[out_cnt])
		out_cnt += 1
		ping, pong = pong, ping

	cboxer.compute(out[pong], result[out_cnt])

print("****Finished****")

#print(power_list)
length = (out_cnt + 1) * BATCH_SIZE
print("# of output: {:d}".format(length))
fps = length/total_time
print("Total time: {:f}s, FPS: {:f}".format(total_time, fps))

power = np.array(power_list)
avgPower = power.mean()
total_energy = avgPower * total_time
print("Total energy: {:f}mJ; Averge Power: {:f}mW".format(total_energy, avgPower))

result = result.reshape(-1, 4)
ground_truth_box = load_ground_truth('ground_truth.txt')
IoU = calculate_aiou(ground_truth_box, result)

f = open("txt/{}.txt".format(args.v),"a+")
# v, f, fps, acc, power, energy
f.write("{}\t{}\t{:2f}\t{:2f}\t{:2f}\t{:2f}\n".format(args.v, args.f, fps, IoU*100, avgPower, total_energy))
f.close()

del img
del cLoader
del cboxer
print()